# -*- coding: utf-8 -*-
# @Author  : 张佳玉
# @Time    : 2023/9/6 10:23
# @Function: EMD

import math
from PIL import Image
import numpy as np


def fill_with_zero(str):
    """
    返回指定长度的字符串，原字符串右对齐，前面填充0。
    @param str:带填补的字符串
    @return:八位长度的字符串
    """
    return str.zfill(8)

def get_key(str1, str2):  # 进制转换，str=0 表示五进制转二进制，str=1 表示二进制转五进制
    s = ""
    if str1 == 0:
        i = int(str2, 5)
        k = bin(i).replace('0b', '')
        s = str(k)
        print("s=", s)
        return s
    elif str1 == 1:
        i = int(str2, 2)
        print("i=", i)
        while i > 0:
            k = i % 5
            i = i // 5
            s += str(k)
        print("五进制:", s[::-1], len(s[::-1]))
        return s[::-1]


def mod(x, y):
    """
    取余
    :param x:
    :param y:
    :return:
    """
    return x % y


def E(n):
    """
    性能指标 嵌入效率E
    :param n: 将秘密信息隐藏进 n 个像素中
    :return: 嵌入效率E
    """
    return ((2 * n + 1) * math.log(2, 2 * n + 1)) / 2 * n


def R(n):
    """
     嵌入率
    :param n: 将秘密信息隐藏进 n 个像素中
    :return: 嵌入率R
    """
    return (math.log(2, 2 * n + 1)) / n


def psnr(origin_image, changed_image):
    """
    性能指标 峰值信噪比 Peak Signal-to-Noise Ratio
    :param origin_image: 原图
    :param changed_image: 信息隐藏后图片
    :return: psnr
    """
    origin_image = np.array(Image.open(origin_image).convert('L')).astype(np.float64)
    changed_image = np.array(Image.open(changed_image).convert('L')).astype(np.float64)
    mse = np.mean((origin_image - changed_image) ** 2)
    if mse == 0:
        return 100
    else:
        return 20 * np.log10(255 / np.sqrt(mse))


def rbase_to_decimal(strr, r):
    """
    r进制转换成十进制
    :param strr:
    :param r:
    :return:
    """
    return int(strr, r)


def decimal_to_rbase(n, r):
    """
    十进制转r进制
    :param n:
    :param r:
    :return:
    """
    num = int(n)
    list_1 = []
    while num:
        list_1.append(mod(num, r))
        num = num // r
    return ''.join([str(x) for x in list_1[::-1]])


# 隐藏函数
def EMD(carrier_image, secret_image, storage_image):
    image = Image.open(carrier_image)
    image_grey = image.convert('L')
    # 获取图片的宽和高
    width = image_grey.size[0]
    print("width:" + str(width))
    height = image_grey.size[1]
    print("height:" + str(height))
    count = 0
    # 获取需要隐藏的信息
    key = decimal_to_rbase(rbase_to_decimal(secret_image, 2), 5)
    # 隐藏信息长度
    key_len = len(key)
    print("len=", key_len)
    for h in range(0, height):
        for w in range(0, width, 2):
            g_1 = image_grey.getpixel((w, h))  # 第i位置的像数值
            g_2 = image_grey.getpixel((w + 1, h))  # 第i+1位置的像素值
            print("w", w, 'w+1', w + 1)
            if count == key_len:
                break
            # 计算f值
            f = (g_1 * 1 + g_2 * 2) % 5
            print("f=", f, "g_1=", g_1, "g_2", g_2, "需要隐藏的key-count=", key[count])
            if key[count] == f:
                count += 1
            else:
                s = int(key[count]) - f % 5
                print("s = ", s)
                if s == 1:
                    g_1 += 1
                elif s == 2:
                    g_2 += 1
                elif s == 3:
                    g_2 -= 1
                elif s == 4:
                    g_1 -= 1
                print("new_g_1=", g_1, "new_g_2", g_2)
                image_grey.putpixel((w + 1, h), g_2)
                image_grey.putpixel((w, h), g_1)
            print("count=", count)

        if count == key_len:
            break

    image_grey.show()
    image_grey.save(storage_image)


def EMD_decode(le, secret_image_test):
    image = Image.open(secret_image_test)
    width = image.size[0]
    height = image.size[1]
    str2 = ""
    count = 0
    for h in range(0, height):
        for w in range(0, width, 2):
            if count == le:
                break
            g_1 = image.getpixel((w, h))  # 第i位置的像数值
            g_2 = image.getpixel((w + 1, h))  # 第i+1位置的像素值
            # 计算f值
            f = (g_1 * 1 + g_2 * 2) % 5
            str2 += str(f)
            count += 1
            if count == le:
                print("str2=", str2)
                101011101001
                print("101011101001")
                #print("所求原二进制数据=", decimal_to_rbase(rbase_to_decimal(str2, 5), 2))
                break


if __name__ == '__main__':
    carrier_image_test = "E:/code/2_EMD/image/LenaRGB.bmp"
    storage_image_test = "E:/code/2_EMD/image/lena_with_secret.jpg"
    secret_image_test = "E:/code/2_EMD/image/copyright.jpg"
    EMD(carrier_image_test, '101011101001', storage_image_test)
    print("psnr=", psnr(carrier_image_test, storage_image_test))
    EMD_decode(5, storage_image_test)


